package top.mcwebsite.lib.gradle.plugin.transform

import java.io.File
import java.util.Collections
import java.util.concurrent.Executors
import java.util.concurrent.LinkedBlockingQueue
import java.util.concurrent.ThreadPoolExecutor
import java.util.concurrent.TimeUnit
import java.util.zip.ZipEntry
import java.util.zip.ZipFile
import java.util.zip.ZipInputStream
import kotlinx.coroutines.async
import kotlinx.coroutines.runBlocking
import org.apache.commons.compress.archivers.jar.JarArchiveEntry
import org.apache.commons.compress.archivers.jar.JarArchiveOutputStream
import org.apache.commons.compress.archivers.zip.ParallelScatterZipCreator
import org.apache.commons.compress.archivers.zip.ZipArchiveEntry
import org.apache.commons.compress.parallel.InputStreamSupplier
import org.gradle.api.DefaultTask
import org.gradle.api.file.Directory
import org.gradle.api.file.RegularFile
import org.gradle.api.file.RegularFileProperty
import org.gradle.api.provider.ListProperty
import org.gradle.api.tasks.InputFiles
import org.gradle.api.tasks.Internal
import org.gradle.api.tasks.OutputFiles
import org.gradle.api.tasks.TaskAction
import top.mcwebsite.lib.transform.api.TransformContext
import top.mcwebsite.lib.transform.api.Transformer

abstract class TransformTask : DefaultTask() {

    /**
     * 所有输入的 jar 文件
     */
    @get:InputFiles
    abstract val allJars: ListProperty<RegularFile>

    /**
     * 所有输入的 classes 目录
     */
    @get:InputFiles
    abstract val allDirectories: ListProperty<Directory>

    /**
     * 输出的 mapping 文件
     */
    @get:OutputFiles
    abstract val outputJar: RegularFileProperty

    @get:Internal
    abstract var transformers: Collection<Transformer>

    private val compileClassPath get() = allJars.get().map { it.asFile } + allDirectories.get().map { it.asFile }

    private val entries = Collections.synchronizedSet(mutableSetOf<String>())

    @TaskAction
    fun transformTaskAction() = runBlocking {
        val context = TransformContext()
        transformers.map {
            async {
                it.onBeforeTransform(context)
            }
        }.forEach {
            it.await()
        }

        compileClassPath
            .takeIf { transformers.any { it.isNeedScan() } }
            ?.flatMap(::parseInput)
            ?.flatMap { (path, byteArraySupplier) ->
                // scan
                transformers
                    .map { transform ->
                        async {
                            transform.scan(path, byteArraySupplier)
                        }
                    }
            }
            ?.forEach {
                it.await()
            }
        JarArchiveOutputStream(outputJar.get().asFile.outputStream().buffered()).use { jaos ->
            val creator = ParallelScatterZipCreator(
                ThreadPoolExecutor(
                    10,
                    10,
                    0L,
                    TimeUnit.MILLISECONDS,
                    LinkedBlockingQueue<Runnable>(),
                    Executors.defaultThreadFactory()
                ) { runnable, _ ->
                    runnable.run()
                }
            )
            compileClassPath.map {
                async {
                    it.transform(creator) { byteArray ->
                        transformers.fold(byteArray) { bytes, transformer ->
                            transformer.transform(bytes)
                        }
                    }
                }
            }.forEach {
                it.await()
            }
            creator.writeTo(jaos)
        }
        transformers.map {
            async {
                it.onAfterTransform(context)
            }
        }.forEach {
            it.await()
        }
    }

    private fun File.transform(
        creator: ParallelScatterZipCreator,
        transformer: (ByteArray) -> ByteArray = { it -> it}
    ) {
        when {
            isDirectory -> {
                val baseURI = this.toURI()
                walk()
                    .filter { it.extension.lowercase() == "class" }
                    .forEach { childFile ->
                    if (childFile.isFile) {
                        val relativePath = baseURI.relativize(childFile.toURI()).path
                        childFile.transform(creator, relativePath, transformer)
                    }
                }
            }
            isFile -> {
                when (extension.lowercase()) {
                    "jar" -> ZipInputStream(this.inputStream()).use {
                        it.transform(creator, ::ZipArchiveEntry, transformer)
                    }
                    "class" -> {
                        transformClass(this, name.substring(1), creator, transformer)
                    }
                    else -> println("Not transform file $path")
                }
            }
        }
    }

    private fun File.transform(creator: ParallelScatterZipCreator, name: String, transformer: (ByteArray) -> ByteArray) {
        if (entries.contains(name)) {
            return
        }
        this.inputStream().use {
            val inputStreamSupplier = {
                transformer(readBytes()).inputStream()
            }
            val jarArchiveEntry = JarArchiveEntry(name).apply {
                method = JarArchiveEntry.DEFLATED
            }
            creator.addArchiveEntry(jarArchiveEntry, inputStreamSupplier)
        }
        entries.add(name)
    }

    private fun transformClass(classFile: File, name: String, creator: ParallelScatterZipCreator, transformer: (ByteArray) -> ByteArray) {
        if (entries.contains(name)) {
            return
        }
        classFile.inputStream().use {
            val inputStreamSupplier = {
                transformer(classFile.readBytes()).inputStream()
            }
            val jarArchiveEntry = JarArchiveEntry(name).apply {
                method = JarArchiveEntry.DEFLATED
            }
            creator.addArchiveEntry(jarArchiveEntry, inputStreamSupplier)
        }
        entries.add(name)
    }


    private fun ZipInputStream.transform(
        creator: ParallelScatterZipCreator,
        entryFactory: (ZipEntry) -> ZipArchiveEntry = ::ZipArchiveEntry,
        transformer: (ByteArray) -> ByteArray
    ) {
        while (true) {
            val next = nextEntry ?: break
            val entry = next.takeIf(::isValidate) ?: continue
            if (!entries.contains(entry.name)) {
                entries.add(entry.name)
                val zae = entryFactory(entry)
                val data = readBytes()
                val stream = InputStreamSupplier {
                    if (entry.isDirectory) {
                        data.inputStream()
                    } else {
                        transformer(data).inputStream()
                    }
                }
                creator.addArchiveEntry(zae, stream)
            }
        }
    }

    private fun isValidate(entry: ZipEntry) =
        (entry.isDirectory || entry.name.endsWith(".class")) && !entry.name.startsWith("META-INF/")

    private fun parseInput(file: File) = when {
        file.isDirectory -> {
            val base = file.toURI()
            file.walk()
                .filter { it.isFile }
                .map {
                    base.relativize(it.toURI()).normalize().path to it::readBytes
                }
                .toList()
        }
        file.isFile -> when (file.extension.lowercase()) {
            "jar", "zip" -> ZipFile(file).use { zip ->
                zip.entries().asSequence()
                    .filterNot(ZipEntry::isDirectory)
                    .map { entry ->
                        entry.name to { zip.getInputStream(entry).readBytes() }
                    }
                    .toList()
            }
            else -> emptyList()
        }
        else -> emptyList()
    }
}